from __future__ import with_statement
from xml.etree import cElementTree
from cgi import escape
from util import *
import re, cPickle, datetime, shelve


class RawXMLDocParser:
    def parse(self,file):
        '''Parses a raw xml document with cElementTree.'''
        title = ''
        text = ''
        docid = 0
        for event, elem in cElementTree.iterparse(file):
            t = re.sub('{.*}','', elem.tag.lower())
            if t == 'page':
                yield (docid, title, text)
                docid += 1
            elif t == 'title':
                title = elem.text
                #title = tostr(elem.text)
                #title = escape(elem.text,True)  # Tuo le ku zi fang pi.
                #title = escape(tostr1(elem.text),True)
            elif t == 'text':
                text = elem.text
                #text = tostr(elem.text)
                #text = escape(elem.text,True)   # Zai ci tuo le ku zi fang pi.
                #text = escape(tostr1(elem.text),True)
            elem.clear()

class FileManager:
    '''Deals with data input and output. Outputs data in collection to the dictionary and posting files,
    or reads in from those files and return an instance of collection.'''
    
    @classmethod
    def parse(cls,file,index):
        '''Use RawXMLDocParser to parse the file in the parameter.
        @param file: path of the file to be parsed.
        @param index: index to be appended.'''

        parser = RawXMLDocParser()
        for d in parser.parse(file):
            index.appendDocument(d[0],d[1],d[2])
            print d[0]
        return index

    @classmethod
    def extractdocuments(cls, file, folder):
        parser = RawXMLDocParser()
        for d in parser.parse(file):
            with open(folder+'doc%d.txt'%d[0],'w') as f:
                f.write(tostr1(d[1])+'\n'+tostr1(d[2]))

    @classmethod
    def write(cls, index, filename):
        '''Outputs the data stored in index to files.
        @param index: an instance of Index.
        @param filename: common prefix of the file to be written, e.g., "collections/seasons".'''

        d = shelve.open(filename+'_dictionary.txt')
        d['dictionary'] = cPickle.dumps(index.dictionary)
        for i in xrange(26):
            char = chr(97+i)
            d['prefixes_%s'%char] = cPickle.dumps(sorted(index.dictionary.prefixes[char]))
            d['suffixes_%s'%char] = cPickle.dumps(sorted(index.dictionary.suffixes[char]))
        d.close()
        
        d = shelve.open(filename+'_postings.txt')
        d['index'] = cPickle.dumps(index)
        d['normtfs'] = cPickle.dumps(tuple(index.normTFs))
        d['idfs'] = cPickle.dumps(tuple(index.idfs))
        d['titles'] = cPickle.dumps(tuple(index.titles))
        for k,v in index.postings.items():
            d[str(k)] = cPickle.dumps(v)
        for k,v in index.wordpostings.items():
            d['w_%s'%str(k)] = cPickle.dumps(v)
        d.close()
        return
    
    @classmethod
    def read(cls, dictfile, postingfile, pagerankfile=None, labmnb=None, labsvm=None, advanced=True):
        '''Reads all the data from the dictionary and posting files into a new instance of Collection. Return the instance.
        @param filename: name of the file to be read. not sure if we need this parameter or not.
        @return: an instance of collection containing all the data from the files.'''        
        d = shelve.open(postingfile)
        idx = cPickle.loads(d['index'])
        idx._dictfile = dictfile
        idx._postingfile = postingfile
        if advanced:
            idx.normTFs = cPickle.loads(d['normtfs'])
            idx.idfs = cPickle.loads(d['idfs'])
            idx.titles = cPickle.loads(d['titles'])
            if pagerankfile:
                idx._pagerankfile = pagerankfile
                f = open(idx._pagerankfile, 'r')
                for line in f:
                    idx.pagerank.append(float(line[:-1]))
                idx.pagerank = tuple(idx.pagerank)
                f.close()

            if labmnb:
                idx._labmnbfile = labmnb
                f = open(idx._labmnbfile, 'r')
                for line in f:
                    l = line.replace('\n','').replace('\r','')
                    if l:
                        idx.labmnb.append(tuple(map(int,l.split(' '))))
                    else:
                        idx.labmnb.append(())
                idx.labmnb = tuple(idx.labmnb)
                f.close()
                
            if labsvm:
                idx._labsvmfile = labsvm
                f = open(idx._labsvmfile, 'r')
                for line in f:
                    l = line.replace('\n','').replace('\r','')
                    if l:
                        idx.labsvm.append(tuple(map(int,l.split(' '))))
                    else:
                        idx.labsvm.append(())
                idx.labsvm = tuple(idx.labsvm)
                f.close()
        d.close()
        
        d = shelve.open(dictfile)
        idx.dictionary = cPickle.loads(d['dictionary'])
        #if advanced:
        #    idx.dictionary.prefixes = cPickle.loads(d['prefixes'])
        #    idx.dictionary.suffixes = cPickle.loads(d['suffixes'])
        d.close()
        return idx

    @classmethod
    def readposting(cls, postingfile, termid):
        '''Read the posting list of certain termid from postingfile.
        @param postingfile: postings file to read
        @param termid: term's id
        @return: term's posting list'''
        d = shelve.open(postingfile)
        try:
            postings = cPickle.loads(d[str(termid)])
        except:
            postings = None
        d.close()
        return postings

    @classmethod
    def readwordposting(cls, postingfile, wordid):
        d = shelve.open(postingfile)
        try:
            postings = cPickle.loads(d['w_%s'%str(wordid)])
        except:
            postings = None
        d.close()
        return postings

    @classmethod
    def readprefix(cls, dictfile, prefix=''):
        ''' Read all the terms from dictionary file that starts with prefix.
        @param dictfile: file name.
        @param prefix: prefix to read.
        '''
        d = shelve.open(dictfile)
        prefixes = cPickle.loads(d['prefixes_%s'%prefix])
        d.close()
        return prefixes
    
    @classmethod
    def readsuffix(cls, dictfile, suffix=''):
        ''' Read all the terms from dictionary file that ends with suffix.
        @param dictfile: file name.
        @param suffix: suffix to read.
        '''
        d = shelve.open(dictfile)
        suffixes = cPickle.loads(d['suffixes_%s'%suffix])
        d.close()
        return suffixes

    #deprecated
    #@classmethod
    #def readprefixes(cls, dictfile):
    #    '''Read all the prefixes of terms from dictionary file.
    #    @param dictfile: file name.
    #    @return: list of prefixes'''
    #    d = shelve.open(dictfile)
    #    prefixes = cPickle.loads(d['prefixes'])
    #    d.close()
    #    return prefixes

    #deprecated
    #@classmethod
    #def readsuffixes(cls, dictfile):
    #    '''Read all the suffixes of terms from dictionary file.
    #    @param dictfile: file name.
    #    @return: list of suffixes'''
    #    d = shelve.open(dictfile)
    #    suffixes = cPickle.loads(d['suffixes'])
    #    d.close()
    #    return suffixes
    
    #deprecated
    #@classmethod
    #def readdocument(cls, postingfile, docID):
    #    d = shelve.open(postingfile)
    #    doc = cPickle.loads(d['doc%d'%docID])
    #    d.close()
    #    return doc

    #deprecated
    #@classmethod
    #def writeposting(cls, index, postingfile, termid=None):
    #    d = shelve.open(postingfile)
    #    if termid is None:
    #        for k,v in index.postings.items():
    #            k = str(k)
    #            if d.has_key(k):
    #                p = cPickle.loads(d[k])
    #                p.merge(v)
    #                d[k] = cPickle.dumps(p)
    #            else:
    #                d[k] = cPickle.dumps(v)
    #    else:
    #        k = str(termid)
    #        if d.has_key(k):
    #            p = cPickle.loads(d[k])
    #            p.merge(index.postings[termid])
    #            d[k] = cPickle.dumps(p)
    #        else:
    #            d[k] = cPickle.dumps(index.postings[termid])
    #    d.close()

    













        
    